import os
import re

import torch
from typing import List, Optional

from ..datasets.video_pose_data import _load_hand_pose_labels
from .hand import HandModel


def get_all_hand_model(data_root):
    data_sources = ["real"]
    hand_nums = ["hand_hand", "separate_hand"]
    usages = ["training", "testing"]

    all_users = set()
    all_hand_models = []
    # 收集所有需要处理的文件
    for data_source in data_sources:
        source_root = os.path.join(data_root, data_source)
        for hand_num in hand_nums:
            num_root = os.path.join(source_root, hand_num)
            for usage in usages:
                usage_root = os.path.join(num_root, usage)
                for user in os.listdir(usage_root):
                    user_root = os.path.join(usage_root, user)
                    all_users.add(user_root)
    all_users = sorted(list(all_users), key=lambda x: os.path.basename(x))
    all_users = [os.path.join(user, "recording_00.json") for user in all_users]
    for hand_model_file in all_users:
        hand_model = _load_hand_pose_labels(hand_model_file).hand_model
        all_hand_models.append(hand_model)
    all_users = [re.findall(r"user_\d\d", user)[0] for user in all_users]
    return all_hand_models, all_users


def save_hand_model_to_dict(hand_model: HandModel) -> dict:
    saved_dict = {}
    # 获取有默认值的可选字段名
    optional_fields = hand_model._field_defaults.keys()

    for field_name in hand_model._fields:
        value = getattr(hand_model, field_name)

        # 跳过可选字段的 None 值
        if value is None and field_name in optional_fields:
            continue

        # 处理非None的张量
        if isinstance(value, torch.Tensor):
            num_elements = value.numel()
            # 单元素张量转为标量
            if num_elements == 1:
                saved_dict[field_name] = value.item()
            # 多元素张量转为列表
            else:
                saved_dict[field_name] = value.tolist()
        # 非张量直接保存（理论上不会出现）
        else:
            saved_dict[field_name] = value

    return saved_dict


def combine_hand_models(
        hand_models: List[HandModel],
        weights: Optional[torch.Tensor] = None
) -> HandModel:
    """
    组合多个手部模型生成一个新模型，只融合四个关键字段

    参数:
        hand_models: 基础手部模型列表（至少一个）
        weights: 可选权重张量，指定不同模型的贡献权重（长度需与hand_models一致）

    返回:
        融合后的新手部模型
    """
    if not hand_models:
        raise ValueError("至少需要一个基础手部模型")

    # 获取第一个模型作为参考
    ref_model = hand_models[0]
    n_models = len(hand_models)

    # 设置默认权重（均匀分布）
    if weights is None:
        weights = torch.ones(n_models) / n_models
    elif len(weights) != n_models:
        raise ValueError(f"权重长度不匹配: {len(weights)} != {n_models}")

    # 准备融合的字段
    fields_to_fuse = [
        "joint_rotation_axes",
        "joint_rest_positions",
        "landmark_rest_positions",
        "mesh_vertices"
    ]

    # 创建融合后的参数字典，初始值为参考模型的值
    fused_params = {field: getattr(ref_model, field) for field in HandModel._fields}

    # 对每个需要融合的字段进行处理
    for field in fields_to_fuse:
        # 检查参考模型该字段是否可用
        if getattr(ref_model, field) is None:
            # 如果参考模型没有该字段，跳过融合
            continue

        # 收集所有模型的该字段
        param_tensors = []
        for model in hand_models:
            param = getattr(model, field)

            # 验证参数一致性
            if param is None:
                raise ValueError(f"模型缺失字段 '{field}'，所有模型必须包含此字段")
            if param.shape != getattr(ref_model, field).shape:
                raise ValueError(f"模型参数'{field}'形状不一致: "
                                 f"{param.shape} vs {getattr(ref_model, field).shape}")
            param_tensors.append(param)

        # 修正维度问题：确保权重张量可以正确广播
        stacked = torch.stack(param_tensors, dim=0)

        # 创建正确维度的权重张量
        # 在堆叠张量的每个后续维度上添加大小为1的维度
        weight_shape = [n_models] + [1] * (stacked.dim() - 1)
        weight_tensor = weights.view(*weight_shape)

        # 加权融合
        combined = torch.sum(stacked * weight_tensor, dim=0)

        # 特殊处理：旋转轴归一化
        if field == "joint_rotation_axes":
            combined = torch.nn.functional.normalize(combined, dim=-1)

        # 更新融合后的参数
        fused_params[field] = combined

    # 创建并返回融合后的手部模型
    return HandModel(**fused_params)


def combine_hand_models_batch(
        hand_models: List[HandModel],
        weights: torch.Tensor = None  # shape: [batch_size, num_models]
) -> HandModel:
    """
    批量组合多个手部模型生成新模型，支持批量权重输入

    参数:
        hand_models: 基础手部模型列表
        weights: 批量权重张量，形状为[batch_size, num_models]

    返回:
        融合后的批量手部模型
    """
    if not hand_models:
        raise ValueError("至少需要一个基础手部模型")
    n_models = len(hand_models)
    if weights is None:
        weights = torch.ones(n_models) / n_models
        weights = weights.unsqueeze(0)
    if weights.dim() != 2:
        raise ValueError(f"权重张量应为二维，但实际维度为 {weights.dim()}")

    batch_size, num_models = weights.shape

    if num_models != len(hand_models):
        raise ValueError(f"权重列数({num_models})与基础模型数({len(hand_models)})不匹配")

    # 获取第一个模型作为参考
    ref_model = hand_models[0]

    # 准备融合的字段
    fields_to_fuse = [
        "joint_rotation_axes",
        "joint_rest_positions",
        "landmark_rest_positions",
        "mesh_vertices"
    ]

    # 创建融合后的参数字典，初始值为参考模型的值（但需扩展为批量）
    fused_params = {}

    # 处理非融合字段（直接从第一个模型复制并扩展为批量）
    for field in HandModel._fields:
        if field not in fields_to_fuse:
            value = getattr(ref_model, field)
            if value is not None:
                # 扩展为批量
                fused_params[field] = value.unsqueeze(0).expand(
                    batch_size, *value.shape
                ).clone()
            else:
                fused_params[field] = None

    # 对每个需要融合的字段进行处理
    for field in fields_to_fuse:
        # 检查参考模型该字段是否可用
        if getattr(ref_model, field) is None:
            # 如果参考模型没有该字段，跳过融合
            fused_params[field] = None
            continue

        # 收集所有模型的该字段
        param_tensors = []
        for model in hand_models:
            param = getattr(model, field)
            if param is None:
                raise ValueError(f"模型缺失字段 '{field}'")

            # 扩展为批量：[1, ...] -> [batch_size, ...]
            param_expanded = param.unsqueeze(0).expand(
                batch_size, *param.shape
            )
            param_tensors.append(param_expanded)

        # 堆叠张量：从[batch_size, ...]到[num_models, batch_size, ...]
        stacked = torch.stack(param_tensors, dim=0)

        # 调整权重维度：[batch_size, num_models] -> [batch_size, num_models, 1, 1, ...]
        # 添加足够的维度以匹配堆叠张量的形状
        weight_shape = (batch_size, num_models) + (1,) * (stacked.dim() - 2)
        weight_tensor = weights.view(weight_shape)

        # 移动维度使batch_size在前：[num_models, batch_size, ...] -> [batch_size, num_models, ...]
        stacked = stacked.permute(1, 0, *range(2, stacked.dim()))

        # 加权融合：sum(weights * stacked, dim=1)
        combined = torch.sum(stacked * weight_tensor, dim=1)

        # 特殊处理：旋转轴归一化
        if field == "joint_rotation_axes":
            combined = torch.nn.functional.normalize(combined, dim=-1)

        # 更新融合后的参数
        fused_params[field] = combined

    # 创建并返回融合后的手部模型
    return HandModel(**fused_params)


if __name__ == '__main__':
    data_root = r"C:\workspace\umetrack_data\raw_data"

    all_hand_models = get_all_hand_model(data_root)
    unique_hand_models = sorted(set(all_hand_models))

    # 创建映射字典
    hand_model_to_idx = {model: idx for idx, model in enumerate(unique_hand_models)}
    num_classes = len(unique_hand_models)
    print(hand_model_to_idx)
